import os
import sys
import numpy as np

from PIL import Image
from pathlib import Path


if __name__ == "__main__":

    STEP_VIDEO_CONVERT_25FPS = False
    STEP_VIDEO_EXTACT_AUDIO = False

    STEP_VIDEO_EXTRACT_FRAME = False
    STEP_VIDEO_EXTRACT_LADMK = False
    STEP_VIDEO_EXTRACT_3DMM = False

    STEP_AUDIO_EXTACT_FEATURE = False
    STEP_AUDIO_COMBINE_FEATURE = False
    STEP_AUDIO_TO_EXPRESSION = True

    STEP_CHOOSE_BG = False
    STEP_RENDER = False

    
    path_project = Path("/mnt/data/CODE/Audio-driven-TalkingFace-HeadPose")
    path_data = Path("/mnt/data/CODE/Audio-driven-TalkingFace-HeadPose/data")
    path_video = path_data.joinpath('video/20200708-1.mp4')
    
    video_name = path_video.stem
    path_save = path_data.joinpath("processing", video_name)

    path_audio = path_save.joinpath('audio/audio_orignal.wav')
    path_video_fps = path_save.joinpath('video/video_25fps.mp4')
    path_audio.parent.mkdir(parents=True, exist_ok=True)
    path_video_fps.parent.mkdir(parents=True, exist_ok=True)

    ### convert the target video to 25fps
    if STEP_VIDEO_CONVERT_25FPS:
        cmd = f"cd video; "\
              f"ffmpeg  -loglevel error -i {path_video} -r 25 {path_video_fps}; "
        print(f"[xxxx] RUN: \n [xxxx] {cmd}")
        os.system(cmd)

        # cmd = f"cd video; "\
        #       f"ffmpeg -loglevel error -i tmp.mp4 -vf scale=\"720:-1\" {path_video_fps}"
        # print(f"[xxxx] RUN: \n [xxxx] {cmd}")
        # os.system(cmd)

    if STEP_VIDEO_EXTACT_AUDIO:
        cmd = f'ffmpeg -y -i {str(path_video)} -vn -acodec pcm_s16le -ac 1 -ar 16000 {str(path_audio)}'
        print(f"[xxxx] RUN: \n [xxxx] {cmd}")
        os.system(cmd)

    ### extract frames and landmarks by
    #path_video_fps = path_data.joinpath('video/input/31.mp4')
    if STEP_VIDEO_EXTRACT_FRAME:
        cmd = f"cd video; "\
            f"python utils/extract_frames.py --path_video {path_video_fps} --path_save {path_save}"

        print(f"[xxxx] RUN: \n [xxxx] {cmd}")
        os.system(cmd)

    if STEP_VIDEO_EXTRACT_LADMK:
        ipath_frames = path_save.joinpath("frames")

        cmd = f"cd video;" \
              f"python utils/extract_landmarks.py --path_frames {ipath_frames} --path_save {path_save} --override"
        
        print(f"[xxxx] RUN: \n [xxxx] {cmd}")
        os.system(cmd)


    ### extract 3DMM coeffs per frames
    if STEP_VIDEO_EXTRACT_3DMM:
        # generate manifest
        ipath_frames = path_save.joinpath("frames")
        ipath_manifest = path_save.joinpath("manifest.csv")
        cmd = f"cd video; python utils/generate_manifest.py --path_frames {ipath_frames} --path_manifest {ipath_manifest}"
        print(f"[xxxx] RUN: \n [xxxx] {cmd}")
        os.system(cmd)


        print("---------------------\n----reconstruction-----\n---------------")
        cmd = f"cd reconstruction/3rdpart/rgbface;                  \
                python -u src/infer_batch.py                         \
                --dataroot {str(path_save)}                      \
                --path_manifest {str(ipath_manifest)}    \
                --path_save {str(path_save)}                     \
                --name RGB_Learn_Adjbrow_1                            \
                --batch_size 100                              \
                --num_threads 6                               \
                --epoch 2                                    \
                --gpu_ids  0 "

        print(f"[xxxx] RUN: \n [xxxx] {cmd}")
        os.system(cmd)

    ### Audio-to-Expression
    # 提取deepspeech feature
    if STEP_AUDIO_EXTACT_FEATURE:
        cmd = f"cd audio/3rdpart/deepspeech;                           \
                python src/infer.py                                    \
                --path-model checkpoints/librispeech_pretrained_v2.pth \
                --path-data  {path_audio}        \
                --path-save  {path_save}        \
                --cuda                                                 \
                --save                                                 \
                --verbose"
        
        print(f"[xxxx] RUN: \n [xxxx] {cmd}")
        os.system(cmd)

    
    # combine features
    if STEP_AUDIO_COMBINE_FEATURE:
        cmd = f"cd video;                           \
                python utils/combine_coeff.py --path_data {path_save}"
        
        print(f"[xxxx] RUN: \n [xxxx] {cmd}")
        os.system(cmd)


    if STEP_AUDIO_TO_EXPRESSION:
        ipath_opt = path_project.joinpath("audio/experiments/deepspeechnet/test/rgbface-test-deep-mouth-landmark-crop.yaml")
        cmd = f"cd audio;" \
              f"python src/infer_one.py \
                --path_data {path_save} --opt {str(ipath_opt)}"

        print(f"[xxxxx] RUN: \n[xxxxx] {cmd}")
        os.system(cmd)


    ### Background matching
    if STEP_CHOOSE_BG:
        pass

    ### Render to 3DMM
    if STEP_RENDER:
        cmd = f"cd reconstruction; \
                CUDA_VISIBLE_DEVICES=0; \
                python render.py --path_video {path_video_fps} --path_audio {path_audio} --path_save {path_save}"

        print(f"[xxxxx] RUN: \n[xxxxx] {cmd}")
        os.system(cmd)

